package common

import (
	"errors"
	"math"
	"sync"
)

// MessageAssembler 用于处理消息的分片与合并
type MessageAssembler struct {
	mu          sync.Mutex
	state       assemblerState // 当前组装状态（空闲/正在合并分片）
	buffer      []byte         // 合并分片的缓冲区
	firstOpcode Opcode         // 首帧的Opcode（确定消息类型）
	closed      bool           // 是否已关闭
}

// 组装器状态：空闲（可开始新消息）、正在合并分片
type assemblerState int8

const (
	idle       assemblerState = iota // 空闲状态（未处理消息）
	assembling                       // 正在合并分片消息
)

// NewMessageAssembler 创建一个新的消息组装器
func NewMessageAssembler() *MessageAssembler {
	return &MessageAssembler{
		state: idle,
	}
}

// 添加一个分片帧，返回是否已合并为完整消息及完整消息数据（仅当返回true时有效）
func (a *MessageAssembler) addFrame(frame *Frame) (bool, []byte, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	if a.closed {
		return false, nil, ErrAssemblerClosed
	}

	switch a.state {
	case idle:
		// 空闲状态：只能处理首帧（非Continuation）
		if frame.Opcode == OpContinuation {
			return false, nil, ErrInvalidFragment
		}
		// 记录首帧类型，初始化缓冲区
		a.firstOpcode = frame.Opcode
		a.buffer = append(a.buffer, frame.Payload...)
		// 如果是最后一帧，直接返回完整消息
		if frame.Fin {
			data := a.buffer
			// 重置状态为空闲
			a.reset()
			return true, data, nil
		}
		// 否则进入分片合并状态
		a.state = assembling
		return false, nil, nil

	case assembling:
		// 合并状态：只能处理Continuation帧
		if frame.Opcode != OpContinuation {
			return false, nil, ErrInvalidFragment
		}
		// 追加当前分片数据
		a.buffer = append(a.buffer, frame.Payload...)
		// 如果是最后一帧，返回完整消息
		if frame.Fin {
			data := a.buffer
			// 重置状态为空闲
			a.reset()
			return true, data, nil
		}
		// 继续等待后续分片
		return false, nil, nil

	default:
		return false, nil, errors.New("unknown assembler state")
	}
}

// 将单个帧拆分为多个帧（根据MessageMaxFrameSize）
func (a *MessageAssembler) splitFrame(op Opcode, data []byte, isClient bool) []*Frame {
	var frames []*Frame
	dataLen := len(data)
	if dataLen == 0 {
		// 空消息：单帧（Fin=true）
		frame := &Frame{
			Fin:     true,
			Opcode:  op,
			Masked:  isClient,
			Payload: data,
		}
		return []*Frame{frame}
	}

	// 计算分片数量
	frameCount := int(math.Ceil(float64(dataLen) / float64(MessageMaxFrameSize)))
	for i := 0; i < frameCount; i++ {
		start := i * MessageMaxFrameSize
		end := start + MessageMaxFrameSize
		if end > dataLen {
			end = dataLen
		}

		// 首帧用原始 Opcode，后续用 Continuation
		opcode := op
		if i > 0 {
			opcode = OpContinuation
		}
		// 最后一帧Fin=true
		fin := i == frameCount-1

		frame := &Frame{
			Fin:     fin,
			Opcode:  opcode,
			Masked:  isClient,
			Payload: data[start:end],
		}

		frames = append(frames, frame)
	}
	return frames
}

// 重置组装器状态（清空缓冲区，回到空闲）
func (a *MessageAssembler) reset() {
	a.buffer = a.buffer[:0] // 清空缓冲区（复用底层数组，减少分配）
	a.firstOpcode = 0
	a.state = idle
}

// Close 关闭组装器
func (a *MessageAssembler) Close() {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.closed = true
	a.reset()
}
